# https://www.kaggle.com/competitions/titanic 
import pandas as pd 
import joblib 
from sklearn.ensemble import RandomForestClassifier 
from sklearn.model_selection import train_test_split 
from sklearn.metrics import accuracy_score 


# 1. 读取数据 
train = pd.read_csv('./train.csv')
test = pd.read_csv('./test.csv')

# 2. 数据预处理
def preprocess_data(df):
    # 性别数字化 
    df['Sex'] = df['Sex'].map({'male': 0, 'female': 1})  
    # 年龄缺失值填充为中位数
    df['Age'].fillna(df['Age'].median(), inplace=True)
    # 填补缺失船舱等级 
    df['Embarked'].fillna(df['Embarked'].mode()[0], inplace=True)
    # 登船港口编码
    df['Embarked'] = df['Embarked'].map({'S': 0, 'C': 1, 'Q': 2})
    # 填补票价缺失值 
    df['Fare'].fillna(df['Fare'].median(), inplace=True)

    return df 


train_data = preprocess_data(train) 
test_data = preprocess_data(test) 

# 3. 特征选择
features = ['Pclass', 'Sex', 'Age', 'SibSp', 'Parch', 'Fare', 'Embarked']
X = train_data[features]
y = train_data['Survived']
X_test = test_data[features]

# 4. 划分训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42) 

# 5. 训练模型
clf = RandomForestClassifier(n_estimators=100, random_state=42)
clf.fit(X_train, y_train)

# 6. 验证模型
y_pred = clf.predict(X_val)
print("Validation Accuracy:", accuracy_score(y_val, y_pred))

# 7. 测试集预测
y_pred_test = clf.predict(X_test)

# 8. 保存模型
joblib.dump(clf, 'titanic_model.pkl') 
print('Model saved as titanic_model.pkl')


# 9. 生成提交文件
submission = pd.DataFrame({
    'PassengerId': test['PassengerId'],
    'Survived': y_pred_test
})


submission.to_csv('submission.csv', index=False)
print('Submission file saved as submission.csv') 


# from sklearn.linear_model import LogisticRegression   # 新增逻辑回归模型 

# # 训练逻辑回归模型
# log_reg = LogisticRegression(max_iter=200)
# log_reg.fit(X_train, y_train)
# y_pred_log_reg = log_reg.predict(X_val)
# print("Logistic Regression Validation Accuracy:", accuracy_score(y_val, y_pred_log_reg))

